import json
import os
import tempfile

def update_system_messages(jsonl_file_path):
    """
    将JSONL文件中每条消息的system信息统一修改为指定的苏格拉底式内容
    确保system消息总是第一条消息
    """
    
    # 固定的system消息内容
    new_system_message = {
        "role": "system",
        "content": "你是古希腊哲学家苏格拉底。你的核心教学方法是\"苏格拉底式诘问法\"（Socratic method）。\n请你通过不断提问的方式来引导我思考，帮助我发现自己观点中的逻辑矛盾，并逐步接近问题的本质。\n不要直接给我答案，而是用问题来启发我。"
    }
    
    # 创建临时文件
    temp_file = tempfile.NamedTemporaryFile(mode='w', delete=False, encoding='utf-8')
    
    try:
        with open(jsonl_file_path, 'r', encoding='utf-8') as file, temp_file:
            for line_num, line in enumerate(file, 1):
                line = line.strip()
                if not line:
                    continue
                    
                try:
                    data = json.loads(line)
                    messages = data.get('messages', [])
                    
                    # 处理消息列表
                    if messages:
                        # 移除所有现有的system消息
                        non_system_messages = [msg for msg in messages if msg.get('role') != 'system']
                        
                        # 在开头插入新的system消息
                        messages = [new_system_message] + non_system_messages
                    
                    # 更新数据并写入临时文件
                    data['messages'] = messages
                    json.dump(data, temp_file, ensure_ascii=False)
                    temp_file.write('\n')
                    
                except json.JSONDecodeError as e:
                    print(f"第{line_num}行JSON解析错误: {e}")
                    # 保留原行
                    temp_file.write(line + '\n')
                except Exception as e:
                    print(f"第{line_num}行处理错误: {e}")
                    # 保留原行
                    temp_file.write(line + '\n')
        
        # 用临时文件替换原文件
        temp_file_path = temp_file.name
        os.replace(temp_file_path, jsonl_file_path)
        print(f"文件 {jsonl_file_path} 已成功更新，所有system消息已被统一修改")
        
    except Exception as e:
        print(f"处理文件时发生错误: {e}")
        # 清理临时文件
        if 'temp_file_path' in locals():
            os.unlink(temp_file_path)

# 使用示例
if __name__ == "__main__":
    # 替换为你的JSONL文件路径
    file_path = "../柏拉图.jsonl"
    update_system_messages(file_path)